import pandas as pd
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score
import numpy as np

close_fib_file = "D:\\stoc\\lz_strategy\\fib_strategy\\fib_close_v3.csv"
y_column = "return5"
y_scale = 0.1
# 特征值
tezheng = [
    "quote_rate", "amp_rate", "turnover_rate",
    # "close_ma5", "close_ma20",  ## 不是比例值
    # "turnover_r5", "turnover_r20",

    "close_fib_d1", "close_fib_d2", "close_fib_d3",
    "close_fib_w1", "close_fib_w2", "close_fib_w3",
    "close_fib_m1", "close_fib_m2", "close_fib_m3",

    "turnover_fib_d1", "turnover_fib_d2", "turnover_fib_d3",
    "turnover_fib_w1", "turnover_fib_w2", "turnover_fib_w3",
    "turnover_fib_m1", "turnover_fib_m2", "turnover_fib_m3"
]
"""
特征值放大10倍取整
3日、5日收益
准确率 0.8
收益率 20%左右
todo 打印出时间和股票来具体分析
"""


def buy_flow(x_test, y_test, dec):
    i = 0
    money = 50000
    err_count = 0
    count_map = {}
    err_count_map = {}
    sum_all = 0
    all_count = 0

    for res in dec.predict(x_test):
        if str(res) in count_map:
            count_map[str(res)] = count_map[str(res)] + 1
        else:
            count_map[str(res)] = 1
        if res > 0:
            all_count = all_count + 1
            print(res, "|||", y_test[i])
            if y_test[i] < 0:
                err_count = err_count + 1
                if str(res) in err_count_map:
                    err_count_map[str(res)].append(y_test[i])
                else:
                    err_count_map[str(res)] = [y_test[i]]
            # money = money * (1 + y_test[i] * y_scale/100)
            # money = money + 10000 * (y_test[i] * y_scale/100)
            sum_all = sum_all + y_test[i] * y_scale
            money = money * (1 + y_test[i] * y_scale / 100)
        i = i + 1
    print(count_map)
    print("error", err_count, err_count_map)
    print("--交易了----", all_count, ":", money, "mean", sum_all / all_count)


def buy_flow_v2(df, dec):
    i = 0
    money = 50000
    err_count = 0
    count_map = {}
    err_count_map = {}
    sum_all = 0
    all_count = 0
    y_test = df[y_column].values
    for res in dec.predict(df[tezheng].values):
        if str(res) in count_map:
            count_map[str(res)] = count_map[str(res)] + 1
        else:
            count_map[str(res)] = 1
        if res > 0:
            all_count = all_count + 1
            print(res, "|||", y_test[i], df.iloc[i][y_column], df.iloc[i]["date"], df.iloc[i]["symbol"])
            if y_test[i] < 0:
                err_count = err_count + 1
                if str(res) in err_count_map:
                    err_count_map[str(res)].append(y_test[i])
                else:
                    err_count_map[str(res)] = [y_test[i]]
            sum_all = sum_all + y_test[i]
            money = money * (1 + y_test[i])
        i = i + 1
    print(count_map)
    print("error", err_count, err_count_map)
    print("--交易了----", all_count, ":", money, "mean", sum_all / all_count)


def decision():
    df = pd.read_csv(close_fib_file)
    print(df.columns)
    ### 删除NAN，看下数据的正确性
    df = df.dropna(axis=0, how='any')
    x = df[tezheng].values
    y = np.around((df[y_column] / y_scale).values, 0)
    # x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.5)
    x_train, x_test = x[: 400], x[400:]
    y_train, y_test = y[: 400], y[400:]
    print("X_train_shape:", x_train.shape, " y_train_shape:", y_train.shape)
    print("X_test_shape:", x_test.shape, "  y_test_shape:", y_test.shape)

    from collections import Counter

    print('Counter(data)\n', Counter(np.around(y_test)))
    # max_depth 5 或 6
    # 剪纸
    # 随机森林
    # dec = DecisionTreeClassifier(max_depth=6)
    # dec.fit(x_train, y_train.astype('int'))
    dec = XGBClassifier()
    dec.fit(x_train, np.around(y_train.astype('int')))
    # dec.save_model("stcok_xgb_1.1.model")
    # dec.load_model("stcok_xgb_1.model")
    # print("----", dec.score(x_train, y_train.astype('int')))
    print("----", dec.score(x_test, np.around(y_test).astype('int')))

    y_pred = dec.predict(x_test)
    accuracy = accuracy_score(np.around(y_test.astype('int')), y_pred)
    print("accuarcy: %.2f%%" % (accuracy * 100.0))
    # buy_flow(x_test, y_test, dec)
    buy_flow_v2(df[375:], dec)
    print("------------------end")

    import matplotlib.pyplot as plt

    from xgboost import plot_importance

    fig, ax = plt.subplots(figsize=(10, 15))
    plot_importance(dec, height=0.5, max_num_features=64, ax=ax)
    plt.show()
    # 参数重要性
    important = dict(zip(tezheng, dec.feature_importances_))
    d_order = sorted(important.items(), key=lambda x: x[1], reverse=False)
    for i in d_order:
        print(i)
    print("end")

    # export_graphviz(dec, feature_names=tezheng, out_file="./tree.dot")
    # dot -Tpdf tree.dot -o tree.pdf
# def load_decision():
#     dec = XGBClassifier()
#     dec.load_model("stcok_xgb.model")


if __name__ == '__main__':
    decision()
